// Copyright 2016 The Gosl Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package fftw wraps the FFTW library to perform Fourier Transforms
// using the "fast" method by Cooley and Tukey
package utils

/*
#include <fftw3.h>
*/
/*
#cgo LDFLAGS: -lfftw3 -lm
*/
import "C"

import (
	"fmt"
	"sync"
	"unsafe"

	"github.com/cpmech/gosl/fun/fftw"
)

// Plan1d implements the FFTW3 plan structure; i.e. a "plan" to compute direct or inverse 1D FTs
//
//   Computes:
//                      N-1         -i 2 π j k / N                 __
//     forward:  X[k] =  Σ  x[j] ⋅ e                     with i = √-1
//                      j=0
//
//                      N-1         +i 2 π j k / N
//     inverse:  Y[k] =  Σ  y[j] ⋅ e                     thus x[k] = Y[k] / N
//                      j=0
//
//   NOTE: FFTW says "so you should initialize your input data after creating the plan."
//         Therefore, the plan can be created and reused several times.
//         [http://www.fftw.org/fftw3_doc/Planner-Flags.html]
//         Also: "The plan can be reused as many times as needed. In typical high-performance
//         applications, many transforms of the same size are computed"
//         [http://www.fftw.org/fftw3_doc/Introduction.html]
//
//         Create a new Plan1d with NewPlan1d(...) AND deallocate memory with Free()
//
type Plan1d struct {
	p    C.fftw_plan  // FFTW "plan" structure
	data []complex128 // input
}

// NewPlan1d allocates a new "plan" to compute 1D Fourier Transforms
//
//   data    -- [modified] data is a complex array of length N.
//   inverse -- will perform inverse transform; otherwise will perform direct
//              Note: both transforms are non-normalised;
//              i.e. the user will have to multiply by (1/n) if computing inverse transforms
//   measure -- use the FFTW_MEASURE flag for better optimisation analysis (slower initialisation times)
//              Note: using this flag with given "data" as input will cause the allocation
//              of a temporary array and the execution of two copy commands with size len(data)
//
//   NOTE: (1) the user must remember to call Free to deallocate FFTW data
//         (2) data will be overwritten
//
func NewPlan1d(data []complex128, inverse, measure bool) (o *Plan1d) {

	// allocate new object
	o = new(Plan1d)
	o.data = data

	// set flags
	var sign C.int = C.FFTW_FORWARD
	var flag C.uint = C.FFTW_ESTIMATE
	if inverse {
		sign = C.FFTW_BACKWARD
	}
	if measure {
		flag = C.FFTW_MEASURE
	}

	// the measure flag will change the input; thus a temporary is required
	N := len(data)
	var temp []complex128
	if measure {
		temp = make([]complex128, N)
		copy(temp, data)
	}

	// set FFTW plan
	d := (*C.fftw_complex)(unsafe.Pointer(&o.data[0]))
	o.p = C.fftw_plan_dft_1d(C.int(N), d, d, sign, flag)

	// fix data (changed by 'measure')
	if measure {
		copy(data, temp)
	}
	return
}

// Free frees internal FFTW data
func (o *Plan1d) Free() {
	if o.p != nil {
		C.fftw_destroy_plan(o.p)
	}
}

// Execute performs the Fourier transform
func (o *Plan1d) Execute() {
	C.fftw_execute(o.p)
}

type FFTPlan struct {
	Size int
	Buf  []complex128
	plan *fftw.Plan1d
}

func NewIFFT(s int) *FFTPlan {
	Buf := make([]complex128, s)
	return &FFTPlan{
		Size: s,
		Buf:  Buf,
		plan: fftw.NewPlan1d(Buf, true, true),
	}
}
func NewFFT(s int) *FFTPlan {
	Buf := make([]complex128, s)
	return &FFTPlan{
		Size: s,
		Buf:  Buf,
		plan: fftw.NewPlan1d(Buf, false, true),
	}
}
func (i *FFTPlan) Execute() { i.plan.Execute() }
func (i *FFTPlan) ExecuteAndScale(s float64) {
	i.plan.Execute()
	for n := 0; n < i.Size; n++ {
		i.Buf[n] *= complex(s, 0.)
	}
}
func (i *FFTPlan) Free() { i.plan.Free() }

type FFTQueue struct {
	queue []*FFTPlan
	lock  sync.RWMutex
}

func (c *FFTQueue) Enqueue(i *FFTPlan) {
	c.lock.Lock()
	defer c.lock.Unlock()
	c.queue = append(c.queue, i)
}

func (c *FFTQueue) Dequeue() (r *FFTPlan, err error) {
	if len(c.queue) > 0 {
		c.lock.Lock()
		defer c.lock.Unlock()
		r, c.queue = c.queue[0], c.queue[1:]
		return
	}
	err = fmt.Errorf("Pop Error: Queue is empty")
	return
}

func (c *FFTQueue) Front() (*FFTPlan, error) {
	if len(c.queue) > 0 {
		c.lock.Lock()
		defer c.lock.Unlock()
		return c.queue[0], nil
	}
	return nil, fmt.Errorf("Peep Error: Queue is empty")
}

func (c *FFTQueue) Size() int {
	return len(c.queue)
}

func (c *FFTQueue) Empty() bool {
	return len(c.queue) == 0
}

type FFTPool struct {
	q *FFTQueue
}

func NewFFTPool(size int, f func() *FFTPlan) (r FFTPool) {
	r = FFTPool{&FFTQueue{}}
	for i := 0; i < size; i++ {
		r.q.Enqueue(f())
	}
	return
}

func (r *FFTPool) Get() (*FFTPlan, error) {
	return r.q.Dequeue()
}

func (r *FFTPool) Return(i *FFTPlan) int {
	r.q.Enqueue(i)
	return r.q.Size()
}
